import torch.nn as nn
import torch.nn.functional as F

from ...losses import BuildLoss
from ..base import BaseSegmentor
from ...backbones import BuildNormalization, BuildActivation
from .head import SCTHead, AU_SCTHead,VitGuidanceHead


'''SCT'''
class SCT(BaseSegmentor):
    def __init__(self, cfg, mode):
        super(SCT, self).__init__(cfg, mode)
        align_corners, norm_cfg, act_cfg, head_cfg = self.align_corners, self.norm_cfg, self.act_cfg, cfg['head']
        self.decoder = SCTHead(head_cfg['in_channels'], head_cfg['feats_channels'], cfg['num_classes'], head_cfg['dropout_ratio'], head_cfg['align_corners'])
        # build auxiliary decoder
        self.setauxiliarydecoder(cfg['auxiliary'])
    '''forward'''
    def forward(self, x, targets=None):
        img_size = x.size(2), x.size(3)
        # feed to backbone network
        backbone_outputs = self.backbone_net(x)
        # feed to decoder
        predictions = self.decoder(backbone_outputs[self.cfg['head']['in_index']])
        # forward according to the mode
        if self.mode == 'TRAIN':
            loss, losses_log_dict = self.customizepredsandlosses(
                predictions=predictions, targets=targets, backbone_outputs=backbone_outputs, losses_cfg=self.cfg['losses'], img_size=img_size,auxiliary_cfg=self.cfg['auxiliary']
            )
            return loss, losses_log_dict
        return predictions

    def setauxiliarydecoder(self, auxiliary_cfg):
        norm_cfg, act_cfg, num_classes = self.norm_cfg.copy(), self.act_cfg.copy(), self.cfg['num_classes']
        if auxiliary_cfg is None: return
        if isinstance(auxiliary_cfg, dict):
            auxiliary_cfg = [auxiliary_cfg]
        self.auxiliary_decoder = nn.ModuleList()
        for aux_cfg in auxiliary_cfg:
            if aux_cfg['type'] == "AU_SCTHead":
                dec = AU_SCTHead(aux_cfg['in_channels'], aux_cfg['feats_channels'], self.cfg['num_classes'], aux_cfg['dropout_ratio'], aux_cfg['align_corners'])
            elif aux_cfg['type'] == "VitGuidanceHead":
                dec = VitGuidanceHead(aux_cfg['in_channels'], aux_cfg['feats_channels'], self.cfg['num_classes'], aux_cfg['dropout_ratio'], aux_cfg['align_corners'])
            self.auxiliary_decoder.append(dec)
        if len(self.auxiliary_decoder) == 1:
            self.auxiliary_decoder = self.auxiliary_decoder[0]

    def customizepredsandlosses(self, predictions, targets, backbone_outputs, losses_cfg, img_size, auxiliary_cfg,auto_calc_loss=True, map_preds_to_tgts_dict=None):
        predictions = F.interpolate(predictions, size=img_size, mode='bilinear', align_corners=self.align_corners)
        outputs_dict = {'loss_cls': predictions}
        for idx, (aux_cfg, dec) in enumerate(zip(auxiliary_cfg,self.auxiliary_decoder)):
            backbone_output = backbone_outputs[aux_cfg["in_index"]]
            predictions_aux = dec(backbone_output)
            if aux_cfg['type'] == "AU_SCTHead":
                predictions_aux = F.interpolate(predictions_aux, size=img_size, mode='bilinear', align_corners=self.align_corners)
            outputs_dict[aux_cfg['type']] = predictions_aux
        return self.calculatelosses(predictions=outputs_dict, targets=targets, losses_cfg=losses_cfg, map_preds_to_tgts_dict=map_preds_to_tgts_dict)

    def calculatelosses(self, predictions, targets, losses_cfg, map_preds_to_tgts_dict=None):
        assert len(predictions) == len(losses_cfg), 'length of losses_cfg should be equal to the one of predictions'
        # calculate loss according to losses_cfg
        losses_log_dict = {}
        for loss_name, loss_cfg in losses_cfg.items():
            if map_preds_to_tgts_dict is None:
                losses_log_dict[loss_name] = self.calculateloss(
                    prediction=predictions[loss_name], target=targets['seg_target'], loss_cfg=loss_cfg,
                )
            else:
                losses_log_dict[loss_name] = self.calculateloss(
                    prediction=predictions[loss_name], target=targets[map_preds_to_tgts_dict[loss_name]], loss_cfg=loss_cfg,
                )
        # summarize and convert losses_log_dict
        loss = 0
        for loss_key, loss_value in losses_log_dict.items():
            loss_value = loss_value.mean()
            loss = loss + loss_value
            loss_value = loss_value.data.clone()
            #dist.all_reduce(loss_value.div_(dist.get_world_size())) #分布式训练使用的
            losses_log_dict[loss_key] = loss_value.item()
        losses_log_dict.update({'loss_total': sum(losses_log_dict.values())})
        # return the loss and losses_log_dict
        return loss, losses_log_dict
    '''calculateloss'''
    def calculateloss(self, prediction, target, loss_cfg):
        assert isinstance(loss_cfg, (dict, list))
        # calculate the loss, dict means single-type loss and list represents multiple-type losses
        if isinstance(loss_cfg, dict):
            loss = BuildLoss(loss_cfg)(prediction, target)
        else:
            loss = 0
            for l_cfg in loss_cfg:
                loss = loss + BuildLoss(l_cfg)(prediction, target)
        # return the loss
        return loss
